SageMakerのノートブックインスタンスを時間になったら起動・停止して、アイドル状態が続いたら停止させてみた

SageMakerのノートブックインスタンスを指定した時間になったら自動で起動・停止させて、アイドル状態が一定時間続いた場合にも自動で停止するようにしてみました。

#Amazon SageMaker

#AWS Lambda

#Amazon EventBridge Scheduler

おざわ（じ）

2023.08.18

この記事は公開されてから1年以上経過しています。情報が古い可能性がありますので、ご注意ください。

はじめに

データアナリティクス事業本部のおざわです。在宅勤務で毎日エアコンをつけっぱなしの生活なので今月分の請求が恐ろしいです。。

さて、いつもお世話になっているSageMakerのノートブックインスタンスですが、ときどき起動に時間がかかることがあります。できればあまり待たずに使えたら嬉しいですよね。また、ノートブックインスタンスは、起動している間は課金されますので節約のためにも停止忘れを防ぎたいものです。

ということで、指定した時間になったら自動でインスタンスを起動・停止してみたいと思います。おまけでインスタンスのアイドル状態が一定時間続いたら自動で停止するように設定します。

今回は、こちらのブログを参考にLambdaやインスタンスのタグ設定を行っています。

指定した時間になったらLambdaを実行するためAmazon EventBridge Schedulerを使用します。

一定時間アイドル状態のインスタンスを停止するには、以下のリポジトリで公開されているライフサイクル設定のスクリプトを使用できます。

1. ノートブックインスタンスのタグ設定

ノートブックインスタンスはすでに存在するものとします。インスタンスにタグを設定して、自動起動・停止の対象かどうかを判断させます。今回は、タグのキーとして「auto_start_stop」、値に「true」を設定しています。このタグが存在しなかったり、値がTrueでない場合は無視されます。

2. Lambda関数を作成

このブログにあるコードを少し変更したものを使いました。次に設定するAmazon EventBridge Schedulerからcommandパラメータを受け取って起動と停止の処理を分けています。

import os
import json
import boto3

# 環境変数を読み込む
TAG_KEY = os.environ["TAG_KEY"]
TAG_VALUE = os.environ["TAG_VALUE"]
REGION = os.environ["REGION"]

def lambda_handler(event, context):
    command = event["command"]
    
    status_mapping = {
        "start": "Stopped",
        "stop": "InService"
    }
    
    if command not in status_mapping:
        return {
            "statusCode": 400,
            "body": json.dumps("Invalid command")
        }
    
    current_status = status_mapping[command]
    
    sm = boto3.client("sagemaker", region_name=REGION)
    # インスタンスを取得する
    nb_list = sm.list_notebook_instances(
        StatusEquals=current_status
    )['NotebookInstances']

    for nb in nb_list:
        # インスタンスのタグを取得する
        tags = sm.list_tags(ResourceArn=nb['NotebookInstanceArn'])['Tags']

        # タグから自動停止・起動が有効かどうかを調べる
        tagged_nb = False
        for tag in tags:
            if tag['Key'] == TAG_KEY and tag['Value'] == TAG_VALUE:
                tagged_nb = True
                break

        nb_name = nb["NotebookInstanceName"]
        if tagged_nb and command == "start":
            print(f"starting instance: {nb_name}")
            start_notebook_instance(nb_name, sm)
        elif tagged_nb and command == "stop":
            print(f"stopping instance: {nb_name}")
            stop_notebook_instance(nb_name, sm)

    return {
        "statusCode": 200,
        "body": json.dumps("done")
    }

def stop_notebook_instance(nb_name, sm_client):
    # 指定されたノートブックインスタンスを停止する
    sm_client.stop_notebook_instance(
        NotebookInstanceName=nb_name
    )

def start_notebook_instance(nb_name, sm_client):
    # 指定されたノートブックインスタンスを起動する
    sm_client.start_notebook_instance(
        NotebookInstanceName=nb_name
    )